monai/networks/nets/unet.py (4)

29-33: Add a brief class docstring to the wrapper.
Improves discoverability and meets docstring guidelines.

Apply this diff:

 class _ActivationCheckpointWrapper(nn.Module):
-    def __init__(self, module: nn.Module) -> None:
+    """Apply activation checkpointing to the wrapped module during training."""
+    def __init__(self, module: nn.Module) -> None:
         super().__init__()
         self.module = module

---

134-135: Document the new use_checkpointing arg in the class docstring and user docs.
State trade-offs (memory vs compute), that it’s training-only, incompatible with torch.no_grad, and preserves RNG by default.

Proposed docstring snippet to add under “Args”:

use_checkpointing: if True, apply activation checkpointing to internal sub-blocks
    during training to reduce memory at the cost of extra compute. Checkpointing
    is bypassed in eval and when gradients are disabled. Defaults to False.

I can open a docs patch and add a short example enabling the flag.

---

163-164: Static wrapping caveat: runtime flips won’t take effect.
After init, changing self.use_checkpointing won’t rewrap existing blocks. Either document this or add a small helper to (re)build the model if you expect runtime toggling.

Do you expect users to toggle this at runtime? If yes, I can sketch a safe rewrap helper.

---

210-212: Checkpointing scope is subblock-only; consider an optional broader policy.
Current placement is a good default. If more memory is needed, offer a policy to also wrap down_path/up_path (with a warning about extra compute).

I can add a checkpoint_policy: Literal["subblock","all"] = "subblock" in __init__ and wire it here on request.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Knowledge Base: Disabled due to Reviews > Disable Knowledge Base setting

monai/networks/nets/unet.py (1)

29-43: Good guard + compatibility fallback.
Training/grad-enabled checks and use_reentrant=False with TypeError fallback are the right call. This addresses the prior review note.

monai/networks/nets/unet.py (5)

29-43: Avoid per-iteration TypeError cost: detect use_reentrant support once.
Resolve support at import/init time to prevent raising an exception every forward on older torch.

Apply:

@@
-class _ActivationCheckpointWrapper(nn.Module):
+_SUPPORTS_USE_REENTRANT: bool | None = None
+
+class _ActivationCheckpointWrapper(nn.Module):
@@
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if self.training and torch.is_grad_enabled() and x.requires_grad:
-            try:
-                return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))
-            except TypeError:
-                # Fallback for older PyTorch without `use_reentrant`
-                return cast(torch.Tensor, checkpoint(self.module, x))
-        return cast(torch.Tensor, self.module(x))
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.training and torch.is_grad_enabled() and x.requires_grad:
+            global _SUPPORTS_USE_REENTRANT
+            if _SUPPORTS_USE_REENTRANT is None:
+                try:
+                    # probe once
+                    checkpoint(self.module, x, use_reentrant=False)  # type: ignore[arg-type]
+                    _SUPPORTS_USE_REENTRANT = True
+                except TypeError:
+                    _SUPPORTS_USE_REENTRANT = False
+                except Exception:
+                    # do not change behavior on unexpected errors; fall back below
+                    _SUPPORTS_USE_REENTRANT = False
+            if _SUPPORTS_USE_REENTRANT:
+                return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))
+            return cast(torch.Tensor, checkpoint(self.module, x))
+        return cast(torch.Tensor, self.module(x))

Add outside the hunk (file header):

import inspect  # if you switch to signature probing instead of try/except

Note: PyTorch recommends passing use_reentrant explicitly going forward. (docs.pytorch.org)

---

29-43: TorchScript: make wrapper script-safe.
try/except and dynamic checkpoint calls won’t script. Short-circuit under scripting.

Apply:

@@
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if torch.jit.is_scripting():
+            # Avoid checkpoint in scripted graphs
+            return cast(torch.Tensor, self.module(x))

---

29-43: Docstring completeness.
Add Google-style docstrings for the wrapper’s class/init/forward (inputs, returns, raises).

Example:

@@
-class _ActivationCheckpointWrapper(nn.Module):
-    """Apply activation checkpointing to the wrapped module during training."""
+class _ActivationCheckpointWrapper(nn.Module):
+    """Wrap a module and apply activation checkpointing during training.
+
+    Args:
+        module: The submodule to checkpoint.
+
+    Returns:
+        torch.Tensor: Output tensor from the wrapped submodule.
+
+    Raises:
+        RuntimeError: If checkpoint fails at runtime.
+    """

---

90-92: Tighten the use_checkpointing docstring and add a BN caveat.
Keep it on one Args entry and note the BatchNorm limitation.

-        use_checkpointing: if True, apply activation checkpointing to internal sub-blocks during training to reduce memory
-        at the cost of extra compute. Checkpointing is bypassed in eval and when gradients are disabled. Defaults to False.
+        use_checkpointing: If True, apply activation checkpointing to internal sub-blocks during training to reduce
+            memory at the cost of extra compute. Bypassed in eval and when gradients are disabled. Note: avoid with
+            BatchNorm layers due to running-stat updates during recomputation. Defaults to False.

Reference on RNG determinism (dropout is handled by default). (docs.pytorch.org)

---

217-219: Wrap site is fine; consider BN detection here instead of inside the wrapper.
Keeps policy close to where wrapping occurs and avoids per-instance warnings.

-        if self.use_checkpointing:
-            subblock = _ActivationCheckpointWrapper(subblock)
+        if self.use_checkpointing:
+            has_bn = any(isinstance(m, nn.modules.batchnorm._BatchNorm) for m in subblock.modules())
+            if has_bn:
+                warnings.warn(
+                    "Skipping activation checkpointing for this subblock (contains BatchNorm).",
+                    RuntimeWarning,
+                )
+            else:
+                subblock = _ActivationCheckpointWrapper(subblock)
         return nn.Sequential(down_path, SkipConnection(subblock), up_path)

Rationale: avoids double-updating BN stats during recomputation. (github.com)

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Knowledge Base: Disabled due to Reviews > Disable Knowledge Base setting

monai/networks/nets/unet.py (4)

29-51: Warn once for BN, and add param/return docstring per guidelines.

Avoid repeated RuntimeWarnings and document the wrapper’s contract.

 class _ActivationCheckpointWrapper(nn.Module):
-    """Apply activation checkpointing to the wrapped module during training."""
+    """Apply activation checkpointing to the wrapped module during training.
+
+    Args:
+        module: submodule to wrap.
+    Returns:
+        torch.Tensor: output of the wrapped module.
+    Warnings:
+        Skips checkpointing and emits a RuntimeWarning if the submodule contains
+        BatchNorm to avoid double-updating running stats.
+    """
     def __init__(self, module: nn.Module) -> None:
         super().__init__()
         # Pre-detect BatchNorm presence for fast path
         self._has_bn = any(isinstance(m, nn.modules.batchnorm._BatchNorm) for m in module.modules())
         self.module = module
+        self._bn_warned = False
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self.training and torch.is_grad_enabled() and x.requires_grad:
             if self._has_bn:
-                warnings.warn(
-                    "Activation checkpointing skipped for a subblock containing BatchNorm to avoid double-updating "
-                    "running statistics during recomputation.",
-                    RuntimeWarning,
-                )
+                if not self._bn_warned:
+                    warnings.warn(
+                        "Activation checkpointing skipped for a subblock containing BatchNorm to avoid double-updating "
+                        "running statistics during recomputation.",
+                        RuntimeWarning,
+                    )
+                    self._bn_warned = True
                 return cast(torch.Tensor, self.module(x))
             try:
                 return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))
             except TypeError:
                 # Fallback for older PyTorch without `use_reentrant`
                 return cast(torch.Tensor, checkpoint(self.module, x))
         return cast(torch.Tensor, self.module(x))

Minimal tests to add:

• Training vs eval parity (values match with/without checkpointing).
• BN subblock emits RuntimeWarning and bypasses checkpointing.
• Guard under torch.no_grad() and when requires_grad=False.

---

99-101: Clarify use_checkpointing behavior in docs (BN, grad, training-only, build-time).

Make expectations explicit for users.

-        use_checkpointing: if True, apply activation checkpointing to internal sub-blocks during training to reduce memory
-        at the cost of extra compute. Checkpointing is bypassed in eval and when gradients are disabled. Defaults to False.
+        use_checkpointing: if True, apply activation checkpointing to internal sub-blocks during training to reduce
+            memory at the cost of extra compute. Checkpointing is active only when `self.training` is True, gradients
+            are enabled, and inputs require gradients; it is bypassed in eval and when grads are disabled.
+            Sub-blocks containing BatchNorm are not checkpointed to avoid double-updating running stats.
+            Note: this is a build-time option; changing it after initialization will not retroactively wrap existing
+            sub-blocks. Defaults to False.

---

179-180: Flag is build-time only; consider asserting or documenting.

Changing self.use_checkpointing post-init has no effect since wrappers are created during construction. The doc update above covers this; alternatively, convert to a read-only attribute to avoid confusion.

---

226-228: Wrapping only the subblock is OK; consider optional coverage toggle.

If desired later, expose an opt-in to also wrap down/up blocks for additional savings.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Knowledge Base: Disabled due to Reviews > Disable Knowledge Base setting

monai/networks/nets/unet.py (1)

29-43: Skip checkpointing for subblocks with BatchNorm to avoid double-updating running stats.
Checkpoint recompute updates BN running stats twice under training. Detect BN in the wrapped module and bypass checkpointing with a warning.

Apply this diff:

 class _ActivationCheckpointWrapper(nn.Module):
-    """Apply activation checkpointing to the wrapped module during training."""
+    """Apply activation checkpointing to the wrapped module during training.

+    Skips checkpointing for submodules containing BatchNorm to avoid double-updating
+    running statistics during recomputation.
+    """
     def __init__(self, module: nn.Module) -> None:
         super().__init__()
         self.module = module
+        self._has_bn = any(isinstance(m, nn.modules.batchnorm._BatchNorm) for m in module.modules())

     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if self.training and torch.is_grad_enabled() and x.requires_grad:
+        if self.training and torch.is_grad_enabled() and x.requires_grad:
+            if self._has_bn:
+                warnings.warn(
+                    "Activation checkpointing skipped for a subblock containing BatchNorm "
+                    "to avoid double-updating running statistics during recomputation.",
+                    RuntimeWarning,
+                )
+                return cast(torch.Tensor, self.module(x))
             try:
                 return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))
             except TypeError:
                 # Fallback for older PyTorch without `use_reentrant`
                 return cast(torch.Tensor, checkpoint(self.module, x))
         return cast(torch.Tensor, self.module(x))

monai/networks/nets/unet.py (3)

90-92: Clarify arg docs and surface BN caveat.
Tighten wording and document BN behavior for transparency.

Apply this diff:

-        use_checkpointing: if True, apply activation checkpointing to internal sub-blocks during training to reduce memory
-        at the cost of extra compute. Checkpointing is bypassed in eval and when gradients are disabled. Defaults to False.
+        use_checkpointing: If True, applies activation checkpointing to internal sub-blocks during training to reduce
+            memory at the cost of extra compute. Bypassed in eval mode and when gradients are disabled.
+            Note: sub-blocks containing BatchNorm are executed without checkpointing to avoid double-updating
+            running statistics. Defaults to False.

---

217-219: Placement of wrapper is sensible; consider optional breadth control.
Future enhancement: expose a knob to checkpoint down/up paths too for deeper memory savings on very deep nets.

---

141-142: Add tests to lock behavior.

• Parity: forward/backward equivalence (outputs/grad norms) with vs. without checkpointing.
• Modes: train vs. eval; torch.no_grad().
• Norms: with InstanceNorm and with BatchNorm (assert BN path skips with warning).

I can draft unit tests targeting UNet’s smallest config to keep runtime minimal—want me to open a follow-up?

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Knowledge Base: Disabled due to Reviews > Disable Knowledge Base setting

monai/networks/nets/unet.py (2)

35-36: Missing training and gradient guards causes eval overhead and no_grad crashes.

The forward unconditionally calls checkpoint. This will:

1. Apply checkpointing during inference (eval mode) → unnecessary compute overhead.
2. Fail under torch.no_grad() → runtime error.

Apply this diff:

 def forward(self, x: torch.Tensor) -> torch.Tensor:
-    return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))
+    if self.training and torch.is_grad_enabled() and x.requires_grad:
+        return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))
+    return cast(torch.Tensor, self.module(x))

---

29-37: BatchNorm in checkpointed subblocks will double-update running stats.

Checkpoint recomputes the forward pass during backward, causing BatchNorm layers to update running_mean/running_var twice per training step, skewing statistics.

Consider detecting BatchNorm in __init__ and either warning or skipping checkpoint:

 class _ActivationCheckpointWrapper(nn.Module):
     """Apply activation checkpointing to the wrapped module during training."""
     def __init__(self, module: nn.Module) -> None:
         super().__init__()
         self.module = module
+        self._has_bn = any(isinstance(m, nn.modules.batchnorm._BatchNorm) for m in module.modules())
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))
+        if self.training and torch.is_grad_enabled() and x.requires_grad:
+            if self._has_bn:
+                warnings.warn(
+                    "Activation checkpointing skipped for subblock with BatchNorm to avoid double-update of running stats.",
+                    RuntimeWarning,
+                )
+                return cast(torch.Tensor, self.module(x))
+            return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))
+        return cast(torch.Tensor, self.module(x))

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

monai/networks/nets/unet.py (2)

26-26: Export CheckpointUNet via __all__.

Public class isn’t exported; importing it from monai.networks.nets will fail. Add it to __all__.

-__all__ = ["UNet", "Unet"]
+__all__ = ["UNet", "Unet", "CheckpointUNet"]

---

36-37: Guard checkpoint under grad-enabled training only.

Forward unconditionally checkpointing will crash under torch.no_grad() (common in eval) because checkpoint requires gradients. Gate on self.training, torch.is_grad_enabled(), and a grad-requiring input before calling checkpoint.

     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))
+        if self.training and torch.is_grad_enabled() and x.requires_grad:
+            return cast(torch.Tensor, checkpoint(self.module, x, use_reentrant=False))
+        return cast(torch.Tensor, self.module(x))

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

monai/networks/nets/unet.py (1)

302-307: Docstring required for CheckpointUNet.
Add a Google-style class docstring describing the checkpointing behavior, inherited args, and trade-offs per our guidelines.

 class CheckpointUNet(UNet):
+    """
+    UNet variant that wraps internal connection blocks with activation checkpointing.
+
+    See `UNet` for constructor arguments. During training with gradients enabled,
+    intermediate activations inside encoder–decoder connections are recomputed in
+    the backward pass to reduce peak memory usage at the cost of extra compute.
+    """

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

monai/networks/nets/unet.py (1)

25-25: Sort __all__ alphabetically.

-__all__ = ["UNet", "Unet", "CheckpointUNet"]
+__all__ = ["CheckpointUNet", "UNet", "Unet"]

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

monai/networks/nets/unet.py (1)

25-25: Consider sorting __all__ alphabetically.

Static analysis suggests sorting: ["CheckpointUNet", "UNet", "Unet"].

-__all__ = ["UNet", "Unet", "CheckpointUNet"]
+__all__ = ["CheckpointUNet", "UNet", "Unet"]

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

tests/networks/nets/test_checkpointunet.py (1)

167-204: Add tests verifying checkpointing activates during training.

All tests run in eval mode (via eval_mode or test_script_save), so checkpointing is never engaged. Must verify: (a) forward pass matches UNet in eval mode, and (b) checkpointing works during training with gradients.

Add two tests as suggested in the past review:

def test_checkpoint_parity_eval(self):
    """Verify CheckpointUNet matches UNet output in eval mode."""
    torch.manual_seed(0)
    from monai.networks.nets.unet import UNet
    config = {"spatial_dims": 2, "in_channels": 1, "out_channels": 3,
              "channels": (16, 32, 64), "strides": (2, 2), "num_res_units": 1}
    
    unet = UNet(**config).to(device)
    checkpoint_unet = CheckpointUNet(**config).to(device)
    checkpoint_unet.load_state_dict(unet.state_dict())
    
    test_input = torch.randn(2, 1, 32, 32).to(device)
    
    with eval_mode(unet), eval_mode(checkpoint_unet):
        out_unet = unet(test_input)
        out_checkpoint = checkpoint_unet(test_input)
        self.assertTrue(torch.allclose(out_unet, out_checkpoint, atol=1e-6))

def test_checkpoint_engages_training(self):
    """Verify checkpointing activates during training."""
    net = CheckpointUNet(
        spatial_dims=2, in_channels=1, out_channels=3,
        channels=(16, 32, 64), strides=(2, 2), num_res_units=1
    ).to(device)
    net.train()
    test_input = torch.randn(2, 1, 32, 32, requires_grad=True, device=device)
    
    output = net(test_input)
    loss = output.sum()
    loss.backward()
    
    self.assertIsNotNone(test_input.grad)

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

tests/networks/nets/test_checkpointunet.py (3)

81-86: Add Google-style docstring.

Per coding guidelines, test methods require docstrings describing purpose, parameters, and expected behavior.

Apply this diff:

     @parameterized.expand(CASES)
     def test_shape(self, input_param, input_shape, expected_shape):
+        """Verify CheckpointUNet produces expected output shapes.
+
+        Args:
+            input_param: UNet constructor arguments.
+            input_shape: Input tensor dimensions.
+            expected_shape: Expected output tensor dimensions.
+        """
         net = CheckpointUNet(**input_param).to(device)

---

88-93: Add Google-style docstring.

Per coding guidelines, test methods require docstrings.

Apply this diff:

     def test_script(self):
+        """Verify CheckpointUNet is scriptable via TorchScript."""
         net = CheckpointUNet(

---

95-99: Add Google-style docstring.

Per coding guidelines, test methods require docstrings.

Apply this diff:

     def test_ill_input_shape(self):
+        """Verify RuntimeError raised for mismatched input shape."""
         net = CheckpointUNet(spatial_dims=2, in_channels=1, out_channels=3, channels=(16, 32, 64), strides=(2, 2))

tests/networks/nets/test_checkpointunet.py (1)

119-139: Test logic is sound; minor redundancy in assertion.

The test verifies gradient propagation during training, indirectly confirming checkpointing works. Line 139's assertIsNotNone(grad_norm) is redundant since grad_norm is always a tensor after line 134.

Optionally remove the redundant assertion:

         # gradient flow check
         grad_norm = sum(p.grad.abs().sum() for p in net.parameters() if p.grad is not None)
         self.assertGreater(grad_norm.item(), 0.0)
 
         # checkpointing should reduce activation memory use; we can't directly assert memory savings
         # but we can confirm no runtime errors and gradients propagate correctly
-        self.assertIsNotNone(grad_norm)

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

tests/networks/nets/test_checkpointunet.py (1)

125-186: Add Google-style docstrings for tests.
Docstrings are missing or not in Google style for test_shape, test_script, test_checkpointing_equivalence_eval, and test_checkpointing_activates_training. Please add compliant summaries and sections so the tests meet the documented requirement. As per coding guidelines

@@
 class TestCheckpointUNet(unittest.TestCase):
     @parameterized.expand(CASES)
     def test_shape(self, input_param, input_shape, expected_shape):
+        """Validate CheckpointUNet output shapes across configurations.
+
+        Args:
+            input_param: Mapping of constructor kwargs for the network under test.
+            input_shape: Shape tuple for the synthetic input tensor.
+            expected_shape: Expected output tensor shape.
+        """
         net = CheckpointUNet(**input_param).to(device)
         with eval_mode(net):
             result = net.forward(torch.randn(input_shape).to(device))
             self.assertEqual(result.shape, expected_shape)
 
     def test_script(self):
-        """
-        TorchScript doesn't support activation-checkpointing (torch.utils.checkpoint) calls inside the module.
-        To keep the test suite validating TorchScript compatibility, script the plain UNet (which is scriptable),
-        rather than the CheckpointUNet wrapper that uses checkpointing internals.
-        """
+        """Script the baseline UNet to maintain TorchScript coverage.
+
+        Raises:
+            RuntimeError: If scripting the baseline UNet fails.
+        """
         net = UNet(
             spatial_dims=2, in_channels=1, out_channels=3, channels=(16, 32, 64), strides=(2, 2), num_res_units=0
         )
         test_data = torch.randn(16, 1, 32, 32)
         test_script_save(net, test_data)
 
     def test_checkpointing_equivalence_eval(self):
-        """Ensure that CheckpointUNet matches standard UNet in eval mode (checkpointing inactive)."""
+        """Confirm eval parity when checkpointing is inactive.
+
+        Raises:
+            AssertionError: If shapes or outputs diverge beyond tolerances.
+        """
         params = dict(
             spatial_dims=2, in_channels=1, out_channels=2, channels=(8, 16, 32), strides=(2, 2), num_res_units=1
         )
@@
     def test_checkpointing_activates_training(self):
-        """Ensure checkpointing triggers recomputation under training and gradients propagate."""
+        """Verify checkpointing recomputes activations during training.
+
+        Raises:
+            AssertionError: If gradients are not propagated through the model.
+        """
         params = dict(
             spatial_dims=2, in_channels=1, out_channels=1, channels=(8, 16, 32), strides=(2, 2), num_res_units=1
         )

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

tests/networks/nets/test_checkpointunet.py (1)

126-183: Add Google-style docstrings with Args sections.

Per coding guidelines, docstrings must describe each parameter. All test methods have brief one-liners but lack Args: sections.

Example for test_shape:

     @parameterized.expand(CASES)
     def test_shape(self, input_param, input_shape, expected_shape):
-        """Validate CheckpointUNet output shapes across configurations."""
+        """Validate CheckpointUNet output shapes across configurations.
+
+        Args:
+            input_param: Dictionary of UNet constructor arguments.
+            input_shape: Tuple specifying input tensor dimensions.
+            expected_shape: Tuple specifying expected output tensor dimensions.
+        """
         net = CheckpointUNet(**input_param).to(device)

Apply similar updates to test_script, test_checkpointing_equivalence_eval, and test_checkpointing_activates_training.

tests/networks/nets/test_checkpointunet.py (2)

149-153: Fragile weight synchronization; use load_state_dict instead.

Relying on manual_seed before each instantiation is fragile—if initialization order differs between UNet and CheckpointUNet (e.g., wrapper layers init first), weights diverge. Use load_state_dict to ensure identical parameters.

-        torch.manual_seed(42)
-        net_plain = UNet(**params).to(device)
-
-        torch.manual_seed(42)
-        net_ckpt = CheckpointUNet(**params).to(device)
+        torch.manual_seed(42)
+        net_plain = UNet(**params).to(device)
+        net_ckpt = CheckpointUNet(**params).to(device)
+        net_ckpt.load_state_dict(net_plain.state_dict())

---

164-165: Tighten tolerance for identical operations.

A threshold of 1e-3 is loose when both models have identical weights and run in eval mode. Use 1e-6 or torch.allclose with default tolerances.

-        diff = torch.mean(torch.abs(y_ckpt - y_plain)).item()
-        self.assertLess(diff, 1e-3, f"Eval-mode outputs differ more than expected (mean abs diff={diff:.6f})")
+        self.assertTrue(
+            torch.allclose(y_ckpt, y_plain, atol=1e-6, rtol=1e-5),
+            f"Eval-mode outputs differ: max abs diff={torch.max(torch.abs(y_ckpt - y_plain)).item():.6e}"
+        )

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

tests/networks/nets/test_checkpointunet.py (1)

125-130: Consider adding parameter descriptions to docstring.

The parameterized test has three parameters but the docstring doesn't describe them. Adding an Args section would improve clarity per coding guidelines.

-    """Validate CheckpointUNet output shapes across configurations."""
+    """Validate CheckpointUNet output shapes across configurations.
+
+    Args:
+        input_param: Dictionary of UNet constructor arguments.
+        input_shape: Tuple specifying input tensor dimensions.
+        expected_shape: Tuple specifying expected output tensor dimensions.
+    """

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

tests/networks/nets/test_checkpointunet.py (1)

125-130: Consider documenting parameterized test arguments.

The docstring omits the three arguments provided by @parameterized.expand. While the brief docstring is present, adding an Args: section would clarify what input_param, input_shape, and expected_shape represent for future maintainers.

Example:

     @parameterized.expand(CASES)
     def test_shape(self, input_param, input_shape, expected_shape):
-        """Validate CheckpointUNet output shapes across configurations."""
+        """Validate CheckpointUNet output shapes across configurations.
+
+        Args:
+            input_param: Dictionary of UNet constructor arguments.
+            input_shape: Tuple specifying input tensor dimensions.
+            expected_shape: Tuple specifying expected output tensor dimensions.
+        """
         net = CheckpointUNet(**input_param).to(device)

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting

tests/networks/nets/test_checkpointunet.py (1)

144-150: Set seed before creating input tensor.

Input x is created before seeds are set. While functionally correct (same x feeds both networks), setting the seed first ensures full reproducibility.

Apply this diff:

-    x = torch.randn(2, 1, 32, 32, device=device)
-
     torch.manual_seed(42)
+    x = torch.randn(2, 1, 32, 32, device=device)
     net_plain = UNet(**params).to(device)
 
     torch.manual_seed(42)

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Cache: Disabled due to data retention organization setting

Knowledge base: Disabled due to Reviews -> Disable Knowledge Base setting